//#include "stdafx.h"
#include "DumpCatch.h"

#include <Dbghelp.h>
#include <Shlwapi.h>
#include <stdexcept>

#pragma comment(lib, "Shlwapi.lib")
#pragma comment(lib, "Dbghelp.lib")

void purecall_handler(void)
{    
    throw std::invalid_argument("");
}

void invalid_parameter_handler(const wchar_t* expression, 
								const wchar_t* function, 
								const wchar_t* file, 
								unsigned int line, 
								uintptr_t pReserved)
{
    throw std::invalid_argument("");
}

void CDumpCatch::SetInvalidHandle()
{
#if _MSC_VER >= 1400  // MSVC 2005/8
    m_prevIPH = _set_invalid_parameter_handler(invalid_parameter_handler);
#endif  // _MSC_VER >= 1400
	//At application, this call can stop show the error message box.
    m_prevPCH = _set_purecall_handler(purecall_handler);
}

void CDumpCatch::UnSetInvalidHandle()
{
#if _MSC_VER >= 1400  // MSVC 2005/8
    _set_invalid_parameter_handler(m_prevIPH);
#endif  // _MSC_VER >= 1400
	//At application this can stop show the error message box.
    _set_purecall_handler(m_prevPCH); 
}

//LPTOP_LEVEL_EXCEPTION_FILTER WINAPI _SetUnhandledExceptionFilter(LPTOP_LEVEL_EXCEPTION_FILTER)
//{
//    return NULL;
//}


BOOL MiniDumpFile(LPCTSTR strPath, PEXCEPTION_POINTERS pException)
{
	BOOL bResult = FALSE;

	HANDLE hDumpFile = ::CreateFile(
		strPath,
		GENERIC_WRITE,
		0,
		NULL,
		CREATE_ALWAYS,
		FILE_ATTRIBUTE_NORMAL,
		NULL);
	if (hDumpFile != INVALID_HANDLE_VALUE)
	{
		MINIDUMP_EXCEPTION_INFORMATION mei;
		mei.ExceptionPointers = pException;
		mei.ThreadId = ::GetCurrentThreadId();
		mei.ClientPointers = TRUE;

		bResult = ::MiniDumpWriteDump(
					::GetCurrentProcess(),
					::GetCurrentProcessId(),
					hDumpFile,
					MiniDumpWithFullMemory,
					&mei,
					NULL,
					NULL);

		::CloseHandle(hDumpFile);
	}

	return bResult;
}

DCCallback	G_fnCallBack = nullptr;
LPVOID		G_lpParameter = nullptr;

LONG WINAPI MyUnhandledExceptionFilter(PEXCEPTION_POINTERS pException)
{
	if (G_fnCallBack != nullptr && !G_fnCallBack(G_lpParameter))
		return EXCEPTION_CONTINUE_SEARCH;

	TCHAR szDumpFile[MAX_PATH] = { 0 };
	TCHAR szDateTime[MAX_PATH] = { 0 };

	::GetModuleFileName(NULL, szDumpFile, MAX_PATH);
	::PathRemoveExtension(szDumpFile);

	SYSTEMTIME time;
	::GetSystemTime(&time);
	wsprintf(
		szDateTime, 
		TEXT("_%04d%02d%02d%02d%02d%02d.dmp"), 
		time.wYear, time.wMonth,  time.wDay, 
		time.wHour, time.wMinute, time.wSecond);

	lstrcat(szDumpFile, szDateTime);

	if (MiniDumpFile(szDumpFile, pException))
		return EXCEPTION_EXECUTE_HANDLER;
	
	return EXCEPTION_CONTINUE_SEARCH;
}

CDumpCatch::CDumpCatch()
{
	SetInvalidHandle();
	InstallExceptionHandle();
}

CDumpCatch::~CDumpCatch()
{
	UnSetInvalidHandle();
	RemoveExceptionHandle();
}

CDumpCatch & CDumpCatch::SetCallback(DCCallback fnCallBack, LPVOID lpParameter)
{
	G_fnCallBack = fnCallBack;
	G_lpParameter = lpParameter;
	return *this;
}

BOOL CDumpCatch::InstallExceptionHandle()
{
    m_prevExceptionFilter = ::SetUnhandledExceptionFilter(MyUnhandledExceptionFilter);
    HookSetUnhandledExceptionFilter();
    return TRUE;
}

BOOL CDumpCatch::RemoveExceptionHandle()
{
	UnhookSetUnhandledExceptionFilter();
    if(m_prevExceptionFilter != NULL)
    {
        ::SetUnhandledExceptionFilter(m_prevExceptionFilter);
        m_prevExceptionFilter = NULL;
    }
    return TRUE;
}


BOOL CDumpCatch::HookSetUnhandledExceptionFilter()
{
	BOOL bResult = FALSE;
	ZeroMemory(m_btSUEF, sizeof(m_btSUEF));

	HMODULE hKernel32 = GetModuleHandle(TEXT("kernel32.dll"));
	if(hKernel32 == NULL)
		hKernel32 = LoadLibrary(TEXT("kernel32.dll"));

	if (hKernel32 != NULL)
	{
		LPVOID lpSUEF = ::GetProcAddress(hKernel32, "SetUnhandledExceptionFilter");
		if(lpSUEF != NULL)
		{
			memcpy(m_btSUEF, lpSUEF, sizeof(m_btSUEF));

			//BYTE btJump[5] = { 0xE9, 0x00, 0x00, 0x00, 0x00 };
			//*((DWORD *)(btJump + 1)) = (DWORD)_SetUnhandledExceptionFilter - (DWORD)lpSUEF - 5;

#ifdef _WIN64
			//xor eax,eax
			//ret
			BYTE btEmpty[sizeof(m_btSUEF)] = { 0x33, 0xC0, 0xC3 };
#else
			//xor eax,eax
			//ret 0004
			BYTE btEmpty[sizeof(m_btSUEF)] = { 0x31, 0xC0, 0xC2, 0x04, 0x00 };
#endif

			DWORD dwOldProtect;
			if (::VirtualProtect(lpSUEF, sizeof(btEmpty), PAGE_EXECUTE_READWRITE, &dwOldProtect))
			{
				memcpy(lpSUEF, btEmpty, sizeof(btEmpty));
				::VirtualProtect(lpSUEF, sizeof(btEmpty), dwOldProtect, &dwOldProtect);
				bResult = TRUE;
			}
		}

		::FreeLibrary(hKernel32);
	}

	return bResult;
}

BOOL CDumpCatch::UnhookSetUnhandledExceptionFilter()
{
	if (m_btSUEF[0] == 0)
		return FALSE;

	BOOL bResult = FALSE;

	HMODULE hKernel32 = LoadLibrary(TEXT("kernel32.dll"));

	if (hKernel32 != NULL)
	{
		LPVOID lpSUEF = ::GetProcAddress(hKernel32, "SetUnhandledExceptionFilter");
		if (lpSUEF != NULL)
		{
			DWORD dwOldProtect;
			if (::VirtualProtect(lpSUEF, sizeof(m_btSUEF), PAGE_EXECUTE_READWRITE, &dwOldProtect))
			{
				memcpy(lpSUEF, m_btSUEF, sizeof(m_btSUEF));
				::VirtualProtect(lpSUEF, sizeof(m_btSUEF), dwOldProtect, &dwOldProtect);
				bResult = TRUE;
			}
		}

		::FreeLibrary(hKernel32);
	}

	return bResult;
}
